iT邦幫忙

2026 iThome 鐵人賽

DAY 20
0
Software Development

GPU 效能優化實戰:30 天從 Kernel 到 Profiling (重賽版)系列 第 20

Day 20:再做一次比較激進的版本]重賽版[

  • 分享至 

  • xImage
  •  

還是先AI slob 剩下的 我們之後再說
這幾天在參加活動@@
希望過幾天還有存量可以繼續QAQ
可能要有四天的 這個 先等我一下 我後續補給你ㄌ


昨天兩張表已經把位置標出來了。

跟 Dao 比,N=14 我們還輸;N=15 後來贏,而且 diff 是 0;N=16 開始付 2-pass 的 memory 錢。跟 Tensor Core 比,慢的不是 MMA,是餵資料。

今天做 Day 15 預告的最後一件事:如果分析已經說明原本寫法輸在哪裡,就把路徑重新切乾淨。這裡說的激進,不是再塞更多 trick,而是比較像 FlashAttention 那種精神——中間結果能不寫出去,就不要寫出去;kernel 邊界用 profile 反推,不靠感覺切。


切之前:N=15 其實已經在付 2-pass 的錢

昨天表一的 N=15 = 82.98 us,是切完以後的數字。切之前不是這樣。

原本 N=15 會被拆成:

Pass 1:N=10 的 single-block kernel
        每個 thread 只拿 1 個 float4(4 個 float)
Pass 2:剩下 5 層,走 tiled shared memory
        一次只搬 TILE_COLS = 16 條 column

兩個 kernel,中間把整份 state 寫回 global memory 再讀回來。batch=128 時,這一份是:

128 × 32768 × 4 bytes = 16 MB

光這一趟 round-trip 就是多讀寫 16 MB。Pass 2 的 tile 又只有 16 條 column 寬,bandwidth 很難打滿。butterfly 層數跟切完之後一樣是 15 層;帳單幾乎全在搬資料和第二次 launch。

當時量到大約 278 us。Dao 在同一題大約 98 us。我們慢的不是加減法,是太早離開 single-block。

NCU 在 N=14 還給過一個很像結論的暗示:chunk 已經 64 KB,再往 N=15 把整份丟進 shared memory 會直接撞限制。如果只讀這句,下一步會變成「那 N=15 只好 2-pass,去優化 Pass 2」。

那就是切錯邊界。


第一刀:N=15 不要拆 pass,shared memory 只拿來換位置

N=15 是 32768 個 float,128 KB。它不適合全部住在 shared memory,但可以住在 256 個 thread 的 register 裡:

256 threads × 32 個 float4 = 32768 floats
每個 thread 持有 128 個 float

shared memory 大約 32 KB,只做 warp 之間的 corner turning,不當整份 state 的倉庫。這跟 Dao 的切法是同一類:butterfly 在 register / warp shuffle 裡做完,smem 只負責換 layout。

切完之後,同一合約(FP32, batch=128):

N=15 時間 相對 Dao diff
切之前:10+5 兩個 pass ~278 us 約 0.35× 沒對齊
Dao 98.35 us
切之後:一個 block 做完 82.98 us 1.19× 0.0e+00

278 us → 82.98 us,大約 3.3 倍。省掉的是那一趟 16 MB round-trip,加上一次 kernel launch,加上 Pass 2 那條太窄的 tile。diff = 0 則是因為這次連加減順序都跟 Dao 對上了。

occupancy 看起來不會漂亮。128 個 float 放 register,compiler 隨時可能 spill。但表上的 3.3 倍說明:這題裡,多一趟 global memory 比 occupancy 低更傷。要盯的是 local memory load/store 有沒有冒出來,不是看到 occupancy 低就把 kernel 拆回兩個 pass。


第二刀:因為 N=15 吃得下,Pass 2 才能變輕

第一刀不只救 N=15。它改變 large-N 怎麼切。

切之前,第一段只吃到 10
    N=16 → 剩下 6 層(超過 5,必須走 tiled smem)
    N=20 → 剩下 10 層

切之後,第一段吃到 15
    N=16 → 剩下 1 層
    N=20 → 剩下 5 層(剛好一個 warp)

剩下 ≤ 5 層時,Pass 2 可以走 Day 17 那顆 shuffle kernel:每個 lane 拿一個 row,__shfl_xor_sync 做完 column butterfly。block 裡不再走 shared memory 的完整交換。

跨 block 的 global memory 仍要付。2-pass 的下限還在。但 block 內能省的,就不要再付一次。

昨天表一的 N=16 = 584.56 usN=20 = 14.71 ms,就是這刀切完以後的數字。Dao 在這裡已經不支援,所以沒有外部對手。我們能確定的是:584 us 對上大約 134 MB traffic,有效頻寬約 229 GB/s,已經接近「在搬資料」。再往上加 Tensor Core,先問的不該是 MMA 能不能更快,而是還能不能少一個 pass。

如果第一段沒推進到 15,N=16 的 Pass 2 會剩下 6 層,連 warp shuffle 都走不成,還要再付一次 tiled smem 的 bank conflict 和更窄的 tile。那筆 before 我們沒有單獨留成一張乾淨的表;能確定的是切法本身比較重。


第三刀:真的需要 shared memory 時,把 memory path 切乾淨

剩下超過 5 層,shuffle 不夠用。這時才回到 tiled shared memory。切之前 Pass 2 一次只搬 16 條 column;切完後一次 64 條,load/store 用 float4,pitch pad 4 個 float,需要時用 ld.global.nc

這一刀沒有 N=15 那種 3.3 倍的單一數字,因為它是 large-N 的底線,不是主戰場。它要解決的是:既然 global memory 避不掉,就不要再用 16-wide 的 tile 跟不連續 access 自己加罰。Day 17 把 kernel 拆過了;今天只記一句——tile 寬度是 bandwidth 問題,不是美觀問題。

end-to-end 也一樣。Tensor Core 那條路如果只量 MMA、不量 transpose,會得到一張很好看、但不是同一題的速度表。昨天表二的 380 us 對 40 us,差的就是把那些邊算進去之後的結果。


切完以後長這樣

N <= 15
    一個 block 做完          ← 第一刀,N=15 從 278 us 到 83 us

N > 15
    Pass 1:先做完 15 層
    Pass 2:
        剩下 <= 5 層 → warp shuffle     ← 第二刀
        剩下  > 5 層 → 64-wide float4 tile ← 第三刀

不是永遠最佳。但它的每一刀都對得上昨天的表:N=15 贏在少搬一次;N=16 慢在避不掉的那一次;Tensor Core 慢在多出來的那幾次。


Tensor Core 還是不該進來救 N=16

N=16 的 584 us 看起來很痛,直覺會想:那把剩下那一層塞進 MMA 呢?

昨天表二已經回答了。TF32 在更小的 N 就因為 transpose 慢九倍。N=16 的下限是 memory traffic。再加 layout conversion,只會把 229 GB/s 那條路弄得更髒。

Tensor Core 該出現的地方,還是資料載入後可以重複使用、每個 output 有大量累加、layout 轉換能被 MMA 攤掉的題。GEMM 是那種題。FWT 不是。

這跟前面 AdaptiveGEMM 不矛盾。同一張卡,GEMM 願意付 Ozaki 的固定成本,是因為乘加夠重。FWT 每次只做一次加減,同一筆帳付不起。


Day 15 到 Day 20

Hadamard 看起來像矩陣乘法
    → 真正該算的是 FWT butterfly
        → 小 N:資料留在一個 block
        → 大 N:才付 2-pass 的 global memory
            → 改寫成 Tensor Core GEMM,layout 比加減法還貴
                → 跟 Dao 放在同一張表上,N=15 切對才贏
                    → 重切:278 us → 83 us,N=16 仍然撞在搬資料

如果前面 GEMM 系列是在學怎麼把精度補回這張卡真正有吞吐的地方,這六天就是在學另一件事:

先尊重 operator 的資料流,再談硬體加速。

Hadamard 這一段,先收到這句,以及那一張 before / after 為止。
明天開始 會來一個完全不一樣的題目(因為我沒題目了 阿哈哈哈 而且又有點忙 我決定把僅剩的案例拿出來ㄌ 後續 嗯 希望我有時間把原版的文章寫ㄉ好一點 )

(希望我之後有力氣 把這個部分的PR 以及案例成果寫成好的一篇blog 如果成功 你就會在下面看到我的文章鏈結 26/09/09 隨筆 就交給未來ㄉ我ㄌ)


上一篇
Day 19:比較標的、profile,還有為什麼不要只看速度表 )重賽版(
下一篇
Day 21:先把 LCN 說清楚,再談 GPU }重賽版{
系列文
GPU 效能優化實戰:30 天從 Kernel 到 Profiling (重賽版)26
圖片
  熱門推薦
圖片
{{ item.channelVendor }} | {{ item.webinarstarted }} |
{{ formatDate(item.duration) }}
直播中

尚未有邦友留言

立即登入留言